{# ----------------------------------------------------------------------------
 # SymForce - Copyright 2022, Skydio, Inc.
 # This source code is under the Apache 2.0 license found in the LICENSE file.
 # ---------------------------------------------------------------------------- #}
{%- import "../util/util.jinja" as util with context -%}
import typing as T

import numpy

from .ops import {{ camelcase_to_snakecase(cls.__name__) }} as ops


class {{ cls.__name__ }}(object):
    {% if doc %}
    """
    Autogenerated Python implementation of :py:class:`{{ cls.__module__ }}.{{ cls.__qualname__ }}`.

    {% for line in doc.split('\n') %}
    {{ line.rstrip() }}
    {% endfor %}
    """
    {% endif %}

    __slots__ = ['data']

    # This is because of an issue where mypy doesn't recognize attributes defined in __slots__
    # See https://github.com/python/mypy/issues/5941
    if T.TYPE_CHECKING:
        data = []  # type: T.List[float]

    def __init__(self,
        {% for arg, size in storage_order %}
        {{ arg }}{% if not loop.last %}, {% endif %}
        {% endfor %}
    ):
        # type: ({% for arg, size in storage_order %}{%if size == 1 %}float{% else %}T.Union[T.Sequence[float], numpy.ndarray]{% endif %}{% if not loop.last %}, {% endif %}{% endfor %}) -> None
        self.data = []
        {% for arg, size in storage_order %}
        {% if size != 1 %}
        {{ util.flatten_if_ndarray(arg, size) | indent(width=8) }}
        {% endif %}
        {% endfor %}

        {% for arg, size in storage_order %}
        self.data.{%if size==1 %}append{% else %}extend{% endif %}({{ arg }})
        {% endfor %}

    def __repr__(self):
        # type: () -> str
        return '<{} {}>'.format(self.__class__.__name__, self.data)


    # --------------------------------------------------------------------------
    # CameraOps
    # --------------------------------------------------------------------------

    {% for spec in specs["CameraOps"] %}
    {{ util.function_declaration(spec, is_method=True, available_classes=[cls]) }}
        {{ util.print_docstring(spec.docstring) | indent(8) }}

        return ops.CameraOps.{{ util.function_name_and_args(spec) }}
    {% endfor %}

    # --------------------------------------------------------------------------
    # StorageOps concept
    # --------------------------------------------------------------------------

    @staticmethod
    def storage_dim():
        # type: () -> int
        return {{ ops.StorageOps.storage_dim(cls) }}

    def to_storage(self):
        # type: () -> T.List[float]
        return list(self.data)

    @classmethod
    def from_storage(cls, vec):
        # type: (T.Sequence[float]) -> {{ cls.__name__ }}
        instance = cls.__new__(cls)

        if isinstance(vec, list):
            instance.data = vec
        else:
            instance.data = list(vec)

        if len(vec) != cls.storage_dim():
            raise ValueError(
                "{} has storage dim {}, got {}.".format(cls.__name__, cls.storage_dim(), len(vec))
            )

        return instance

    # --------------------------------------------------------------------------
    # LieGroupOps concept
    # --------------------------------------------------------------------------

    @staticmethod
    def tangent_dim():
        # type: () -> int
        return {{ ops.LieGroupOps.tangent_dim(cls) }}

    @classmethod
    def from_tangent(cls, vec, epsilon=1e-8):
        # type: (numpy.ndarray, float) -> {{ cls.__name__ }}
        if len(vec) != cls.tangent_dim():
            raise ValueError(
                "Vector dimension ({}) not equal to tangent space dimension ({}).".format(
                    len(vec), cls.tangent_dim()
                )
            )
        return ops.LieGroupOps.from_tangent(vec, epsilon)

    def to_tangent(self, epsilon=1e-8):
        # type: (float) -> numpy.ndarray
        return ops.LieGroupOps.to_tangent(self, epsilon)

    def retract(self, vec, epsilon=1e-8):
        # type: (numpy.ndarray, float) -> {{ cls.__name__ }}
        if len(vec) != self.tangent_dim():
            raise ValueError(
                "Vector dimension ({}) not equal to tangent space dimension ({}).".format(
                    len(vec), self.tangent_dim()
                )
            )
        return ops.LieGroupOps.retract(self, vec, epsilon)

    def local_coordinates(self, b, epsilon=1e-8):
        # type: ({{ cls.__name__ }}, float) -> numpy.ndarray
        return ops.LieGroupOps.local_coordinates(self, b, epsilon)

    def interpolate(self, b, alpha, epsilon=1e-8):
        # type: ({{ cls.__name__ }}, float, float) -> {{ cls.__name__ }}
        return ops.LieGroupOps.interpolate(self, b, alpha, epsilon)

    # --------------------------------------------------------------------------
    # General Helpers
    # --------------------------------------------------------------------------
    def __eq__(self, other):
        # type: (T.Any) -> bool
        if isinstance(other, {{ cls.__name__ }}):
            return self.data == other.data
        else:
            return False
